// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/try_expr.h>
#include <pollux/expression/scoped_var_setter.h>

namespace kumo::pollux::exec {
    void TryExpr::evalSpecialForm(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        ScopedVarSetter throwOnError(context.mutableThrowOnError(), false);
        ScopedVarSetter captureErrorDetails(
            context.mutableCaptureErrorDetails(), false);

        ScopedThreadSkipErrorDetails skipErrorDetails(true);

        // It's possible with nested TRY expressions that some rows already threw
        // exceptions in earlier expressions that haven't been handled yet. To avoid
        // incorrectly handling them here, store those errors and temporarily reset
        // the errors in context to nullptr, so we only handle errors coming from
        // expressions that are children of this TRY expression.
        // This also prevents this TRY expression from leaking exceptions to the
        // parent TRY expression, so the parent won't incorrectly null out rows that
        // threw exceptions which this expression already handled.
        ScopedVarSetter<EvalErrorsPtr> errorsSetter(context.errorsPtr(), nullptr);

        // Allocate error vector to avoid repeated re-allocations for every failed
        // row.
        context.ensureErrorsVectorSize(rows.end());

        inputs_[0]->eval(rows, context, result);

        nullOutErrors(rows, context, result);
    }

    void TryExpr::evalSpecialFormSimplified(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) {
        ScopedVarSetter throwOnError(context.mutableThrowOnError(), false);
        ScopedVarSetter captureErrorDetails(
            context.mutableCaptureErrorDetails(), false);

        ScopedThreadSkipErrorDetails skipErrorDetails(true);

        // It's possible with nested TRY expressions that some rows already threw
        // exceptions in earlier expressions that haven't been handled yet. To avoid
        // incorrectly handling them here, store those errors and temporarily reset
        // the errors in context to nullptr, so we only handle errors coming from
        // expressions that are children of this TRY expression.
        // This also prevents this TRY expression from leaking exceptions to the
        // parent TRY expression, so the parent won't incorrectly null out rows that
        // threw exceptions which this expression already handled.
        ScopedVarSetter<EvalErrorsPtr> errorsSetter(context.errorsPtr(), nullptr);

        inputs_[0]->evalSimplified(rows, context, result);

        nullOutErrors(rows, context, result);
    }

    namespace {
        // Apply onError methods of registered listeners on every row that encounters
        // errors. The error vector must exist.
        void applyListenersOnError(
            const SelectivityVector &rows,
            const EvalCtx &context) {
            const auto *errors = context.errors();
            POLLUX_CHECK_NOT_NULL(errors);

            vector_size_t numErrors = 0;
            rows.applyToSelected([&](auto row) {
                if (errors->hasErrorAt(row)) {
                    ++numErrors;
                }
            });

            if (numErrors == 0) {
                return;
            }

            exprSetListeners().withRLock([&](auto &listeners) {
                if (!listeners.empty()) {
                    for (auto &listener: listeners) {
                        listener->onError(numErrors, context.execCtx()->queryCtx()->queryId());
                    }
                }
            });
        }
    } // namespace

    void TryExpr::nullOutErrors(
        const SelectivityVector &rows,
        EvalCtx &context,
        VectorPtr &result) const {
        const auto *errors = context.errors();
        if (!errors) {
            return;
        }

        if (!errors->hasError()) {
            return;
        }

        applyListenersOnError(rows, context);

        if (result->is_constant_encoding()) {
            // Since it's constant, if any row is NULL they're all NULL, so check row
            // 0 arbitrarily.
            if (result->is_null_at(0)) {
                // The result is already a NULL constant, so this is a no-op.
                return;
            }

            auto size = result->size();
            POLLUX_DCHECK_GE(size, rows.end());

            auto nulls = allocate_nulls(size, context.pool());
            auto rawNulls = nulls->asMutable<uint64_t>();
            rows.applyToSelected([&](auto row) {
                if (errors->hasErrorAt(row)) {
                    bits::set_null(rawNulls, row, true);
                }
            });

            // Wrap in dictionary indices all pointing to index 0.
            auto indices = allocate_indices(size, context.pool());
            result = BaseVector::wrap_in_dictionary(nulls, indices, size, result);
        } else if (
            result.use_count() == 1 && result->is_nulls_writable() &&
            result->size() >= rows.end()) {
            auto *rawNulls = result->mutable_raw_nulls();
            rows.applyToSelected([&](auto row) {
                if (errors->hasErrorAt(row)) {
                    bits::set_null(rawNulls, row, true);
                }
            });
        } else {
            auto nulls = allocate_nulls(rows.end(), context.pool());
            auto *rawNulls = nulls->asMutable<uint64_t>();
            auto indices = allocate_indices(rows.end(), context.pool());
            auto *rawIndices = indices->asMutable<vector_size_t>();

            rows.applyToSelected([&](auto row) {
                rawIndices[row] = row;
                if (errors->hasErrorAt(row)) {
                    bits::set_null(rawNulls, row, true);
                }
            });

            result = BaseVector::wrap_in_dictionary(nulls, indices, rows.end(), result);
        }
    }

    TypePtr TryCallToSpecialForm::resolveType(
        const std::vector<TypePtr> &argTypes) {
        POLLUX_CHECK_EQ(
            argTypes.size(),
            1,
            "TRY expressions expect exactly 1 argument, received: {}",
            argTypes.size());
        return argTypes[0];
    }

    ExprPtr TryCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<ExprPtr> &&compiledChildren,
        bool /* trackCpuUsage */,
        const core::QueryConfig & /*config*/) {
        POLLUX_CHECK_EQ(
            compiledChildren.size(),
            1,
            "TRY expressions expect exactly 1 argument, received: {}",
            compiledChildren.size());
        return std::make_shared<TryExpr>(type, std::move(compiledChildren[0]));
    }
} // namespace kumo::pollux::exec
